import pandas as pd
from pandas.api.types import is_numeric_dtype
from scipy.stats import ttest_ind, chi2_contingency

def compare_groups(df_response, df_nonresponse, alpha=0.05, save_path=None):
    """
    对比响应组和非响应组的逐变量分布差异
    连续变量 → t检验
    分类变量 → 卡方检验
    
    参数:
        df_response: DataFrame, 响应组（有标签）
        df_nonresponse: DataFrame, 非响应组（无标签）
        alpha: 显著性水平 (默认0.05)
        save_path: 若不为 None，则保存结果为 Excel 文件
    
    返回:
        results_df: 包含变量名、变量类型、p值、是否显著的结果表
    """
    # 合并数据，加上分组标记
    df_response = df_response.copy()
    df_nonresponse = df_nonresponse.copy()
    df_response["group"] = 1
    df_nonresponse["group"] = 0
    df_all = pd.concat([df_response, df_nonresponse], ignore_index=True)

    # 区分连续变量和分类变量
    continuous_vars = [col for col in df_all.columns if is_numeric_dtype(df_all[col]) and col != "group"]
    categorical_vars = [col for col in df_all.columns if not is_numeric_dtype(df_all[col]) and col != "group"]

    results = []

    # 连续变量：t检验
    for col in continuous_vars:
        g1 = df_response[col].dropna()
        g0 = df_nonresponse[col].dropna()
        if len(g1) > 0 and len(g0) > 0:
            stat, p = ttest_ind(g1, g0, equal_var=False)  # Welch’s t-test
            results.append({"variable": col, "type": "continuous", "p_value": p})

    # 分类变量：卡方检验
    for col in categorical_vars:
        contingency = pd.crosstab(df_all[col], df_all["group"])
        if contingency.shape[0] > 1 and contingency.shape[1] > 1:
            chi2, p, dof, exp = chi2_contingency(contingency)
            results.append({"variable": col, "type": "categorical", "p_value": p})

    # 结果表
    results_df = pd.DataFrame(results)
    results_df["significant"] = results_df["p_value"] <= alpha

    # 按 p 值排序
    results_df = results_df.sort_values(by="p_value").reset_index(drop=True)

    # 保存为 csv
    if save_path:
        results_df.to_csv(save_path, index=False)
        print(f"结果已保存到: {save_path}")

    return results_df


if __name__ == "__main__":
    df_response = pd.read_pickle(r'data/filtered_raw/data_new.pkl')
    df_nonresponse = pd.read_pickle(r'data/filtered_raw/data_nonresponse.pkl')
    results_df = compare_groups(df_response, df_nonresponse, save_path=r'analysis/results/response_bias.csv')
    print(results_df)